import os
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split

def load_data(data_dir, n_frames=2000):
    asd_data = []
    td_data = []

    # ASD
    asd_folder = os.path.join(data_dir, 'ASD')
    for filename in os.listdir(asd_folder):
        if filename.endswith('.csv'):
            df = pd.read_csv(os.path.join(asd_folder, filename))
            df = df.head(n_frames)
            if len(df) < n_frames:
                pad_len = n_frames - len(df)
                pad = pd.DataFrame(np.zeros((pad_len, df.shape[1])), columns=df.columns)
                df = pd.concat([df, pad], ignore_index=True)
            asd_data.append(df[['Gaze_X', 'Gaze_Y', 'Expression']].values.flatten())

    # TD
    td_folder = os.path.join(data_dir, 'TD')
    for filename in os.listdir(td_folder):
        if filename.endswith('.csv'):
            df = pd.read_csv(os.path.join(td_folder, filename))
            df = df.head(n_frames)
            if len(df) < n_frames:
                pad_len = n_frames - len(df)
                pad = pd.DataFrame(np.zeros((pad_len, df.shape[1])), columns=df.columns)
                df = pd.concat([df, pad], ignore_index=True)
            td_data.append(df[['Gaze_X', 'Gaze_Y', 'Expression']].values.flatten())

    return asd_data, td_data

def preprocess_data(asd_data, td_data):
    X = np.array(asd_data + td_data)
    y = np.array([1]*len(asd_data) + [0]*len(td_data))
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, stratify=y, random_state=42
    )
    # 保存为DataFrame
    train_df = pd.DataFrame(X_train)
    train_df['Label'] = y_train
    test_df = pd.DataFrame(X_test)
    test_df['Label'] = y_test
    return train_df, test_df

def main(data_dir):
    asd_data, td_data = load_data(data_dir)
    train_data, test_data = preprocess_data(asd_data, td_data)
    train_data.to_csv(os.path.join(data_dir, 'train_data.csv'), index=False)
    test_data.to_csv(os.path.join(data_dir, 'test_data.csv'), index=False)
    print("预处理完成，训练集和测试集已保存。")

if __name__ == "__main__":
    data_directory = r'D:\毕设'
    main(data_directory)